import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

# ============= 数据集探索 ==============
# 读取CIFAR-10数据集
cifar10 = tf.keras.datasets.cifar10
(x_train, y_train), (x_test, y_test) = cifar10.load_data()

print('training data shape:', x_train.shape)
print('training labels shape:', y_train.shape)
print('test data shape:', x_test.shape)
print('test labels shape:', y_test.shape)

# plt.imshow(x_train[6])
# plt.show()
# print(y_train[6])

# 定义标签字段，每个数字所代表的图像类型类别的名称
label_dict = {0: 'airplane', 1: 'automobile', 2: 'bird', 3: 'cat', 4: 'deer',
              5: 'dog', 6: 'frog', 7: 'horse', 8: 'ship', 9: 'truck'}

# 特征数据归一化
x_train = x_train.astype('float32') / 255.0
x_test = x_test.astype('float32') / 255.0

# ============= 模型定义 ===============
# 建立CNN模型
model = tf.keras.models.Sequential([
    # 第1个卷积层
    tf.keras.layers.Conv2D(filters=32,
                           kernel_size=(3, 3),
                           input_shape=(32, 32, 3),
                           activation='relu',
                           padding='same'),
    # 防止过拟合
    tf.keras.layers.Dropout(rate=0.3),
    # 第1个池化层
    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
    # 第2个卷积层
    tf.keras.layers.Conv2D(filters=64,
                           kernel_size=(3, 3),
                           activation='relu',
                           padding='same'),
    # 防止过拟合
    tf.keras.layers.Dropout(rate=0.3),
    # 第2个池化层
    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
    # 平坦层
    tf.keras.layers.Flatten(),
    # 输出层
    tf.keras.layers.Dense(10, activation='softmax')
])
# 输出模型摘要
model.summary()

# ================= 模型训练 ===================
# 设置训练参数
training_epochs = 5  # 训练轮数
batch_size = 100  # 单次训练样本数
# 定义训练模式
model.compile(optimizer='adam',  # 优化器
              loss=tf.keras.losses.sparse_categorical_crossentropy,  # 损失函数
              metrics=['accuracy'])  # 评估指标
# 模型训练
train_history = model.fit(x=x_train, y=y_train,
          validation_split=0.2,
          epochs=training_epochs,
          batch_size=batch_size,
          verbose=2)

# 训练过程指标可视化
def show_train_history(train_history, train_metric, val_metric):
    plt.plot(train_history.history[train_metric])
    plt.plot(train_history.history[val_metric])
    plt.title('Train History')
    plt.ylabel(train_metric)
    plt.xlabel('Epoch')
    plt.legend(['train', 'validation'], loc='upper left')
    plt.show()

# 可视化损失与准确率
show_train_history(train_history, 'loss', 'val_loss')
show_train_history(train_history, 'accuracy', 'val_accuracy')

# ================= 评估模型 ===================
test_loss, test_acc = model.evaluate(x_test, y_test, verbose=2)
print('评估模型：')
print(test_loss, test_acc)

# ================= 应用模型 ===================
print('应用模型')
preds = model.predict_classes(x_test)
print(preds)


# ================= 可视化预测结果 ===================
def plot_images_labels_prediction(images,  # 图像列表
                                  labels,  # 标签列表
                                  preds,  # 预测值列表
                                  index=0,  # 从第index个开始显示
                                  num=10):  # 缺省一次显示10幅
    fig = plt.gcf()  # 获取当前图表，Get Current Figure
    fig.set_size_inches(12, 6)  # 1英寸等于2.54cm
    if num > 10:
        num = 10  # 最多显示10个子图
    for i in range(num):
        ax = plt.subplot(2, 5, i+1)  # 获取当前要处理的子图
        ax.imshow(images[index])  # 显示第index个图像
        title = str(i) + ',' + label_dict[labels[index][0]]  # 构建该图上要显示的title信息
        if len(preds) > 0:
            title += ' => ' + label_dict[preds[index]]

        ax.set_title(title, fontsize=10)  # 显示图上的title信息
        index += 1
    plt.show()


plot_images_labels_prediction(x_test, y_test, preds, 0, 10)
